#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据库优化迁移脚本
添加新字段和索引以提升查询性能
"""

import sys
import os
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from datetime import datetime
import logging

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from config import Config
from database import Base, Question, Submission, SubmissionDetail

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def migrate_database():
    """执行数据库迁移"""
    try:
        # 创建数据库引擎
        engine = create_engine(Config.DATABASE_URL)
        SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
        
        logger.info("开始数据库优化迁移...")
        
        with engine.connect() as conn:
            # 开始事务
            trans = conn.begin()
            
            try:
                # 1. 为Questions表添加新字段
                logger.info("为Questions表添加新字段...")
                
                # 检查字段是否已存在
                result = conn.execute(text("""
                    SELECT COUNT(*) as count FROM pragma_table_info('questions') 
                    WHERE name IN ('difficulty_level', 'topic_category', 'created_time', 'updated_time')
                """))
                existing_fields = result.fetchone()[0]
                
                if existing_fields < 4:
                    # 添加新字段
                    conn.execute(text("ALTER TABLE questions ADD COLUMN difficulty_level VARCHAR"))
                    conn.execute(text("ALTER TABLE questions ADD COLUMN topic_category VARCHAR"))
                    conn.execute(text("ALTER TABLE questions ADD COLUMN created_time DATETIME DEFAULT CURRENT_TIMESTAMP"))
                    conn.execute(text("ALTER TABLE questions ADD COLUMN updated_time DATETIME DEFAULT CURRENT_TIMESTAMP"))
                    logger.info("Questions表新字段添加完成")
                else:
                    logger.info("Questions表字段已存在，跳过添加")
                
                # 2. 为Submissions表添加新字段
                logger.info("为Submissions表添加新字段...")
                
                result = conn.execute(text("""
                    SELECT COUNT(*) as count FROM pragma_table_info('submissions') 
                    WHERE name IN ('total_questions', 'correct_answers', 'completion_time')
                """))
                existing_fields = result.fetchone()[0]
                
                if existing_fields < 3:
                    conn.execute(text("ALTER TABLE submissions ADD COLUMN total_questions INTEGER DEFAULT 0"))
                    conn.execute(text("ALTER TABLE submissions ADD COLUMN correct_answers INTEGER DEFAULT 0"))
                    conn.execute(text("ALTER TABLE submissions ADD COLUMN completion_time INTEGER"))
                    logger.info("Submissions表新字段添加完成")
                else:
                    logger.info("Submissions表字段已存在，跳过添加")
                
                # 3. 为SubmissionDetails表添加新字段
                logger.info("为SubmissionDetails表添加新字段...")
                
                result = conn.execute(text("""
                    SELECT COUNT(*) as count FROM pragma_table_info('submission_details') 
                    WHERE name IN ('is_correct', 'answer_time')
                """))
                existing_fields = result.fetchone()[0]
                
                if existing_fields < 2:
                    conn.execute(text("ALTER TABLE submission_details ADD COLUMN is_correct INTEGER DEFAULT 0"))
                    conn.execute(text("ALTER TABLE submission_details ADD COLUMN answer_time DATETIME DEFAULT CURRENT_TIMESTAMP"))
                    logger.info("SubmissionDetails表新字段添加完成")
                else:
                    logger.info("SubmissionDetails表字段已存在，跳过添加")
                
                # 4. 创建索引以优化查询性能
                logger.info("创建数据库索引...")
                
                # Questions表索引
                indexes = [
                    "CREATE INDEX IF NOT EXISTS idx_questions_knowledge_point ON questions(knowledge_point)",
                    "CREATE INDEX IF NOT EXISTS idx_questions_answer ON questions(answer)",
                    "CREATE INDEX IF NOT EXISTS idx_questions_difficulty ON questions(difficulty_level)",
                    "CREATE INDEX IF NOT EXISTS idx_questions_topic ON questions(topic_category)",
                    "CREATE INDEX IF NOT EXISTS idx_questions_created_time ON questions(created_time)",
                    
                    # Submissions表索引
                    "CREATE INDEX IF NOT EXISTS idx_submissions_student_name ON submissions(student_name)",
                    "CREATE INDEX IF NOT EXISTS idx_submissions_class_id ON submissions(class_id)",
                    "CREATE INDEX IF NOT EXISTS idx_submissions_score ON submissions(score)",
                    "CREATE INDEX IF NOT EXISTS idx_submissions_date ON submissions(submission_date)",
                    "CREATE INDEX IF NOT EXISTS idx_submissions_time ON submissions(submission_time)",
                    
                    # SubmissionDetails表索引
                    "CREATE INDEX IF NOT EXISTS idx_submission_details_submission_id ON submission_details(submission_id)",
                    "CREATE INDEX IF NOT EXISTS idx_submission_details_question_id ON submission_details(question_id)",
                    "CREATE INDEX IF NOT EXISTS idx_submission_details_selected_answer ON submission_details(selected_answer)",
                    "CREATE INDEX IF NOT EXISTS idx_submission_details_is_correct ON submission_details(is_correct)",
                    "CREATE INDEX IF NOT EXISTS idx_submission_details_answer_time ON submission_details(answer_time)",
                    
                    # 复合索引
                    "CREATE INDEX IF NOT EXISTS idx_submissions_class_date ON submissions(class_id, submission_date)",
                    "CREATE INDEX IF NOT EXISTS idx_submission_details_submission_question ON submission_details(submission_id, question_id)"
                ]
                
                for index_sql in indexes:
                    try:
                        conn.execute(text(index_sql))
                    except Exception as e:
                        logger.warning(f"索引创建可能已存在: {e}")
                
                logger.info("数据库索引创建完成")
                
                # 5. 更新现有数据的时间戳
                logger.info("更新现有数据的时间戳...")
                
                current_time = datetime.now().isoformat()
                conn.execute(text(f"""
                    UPDATE questions 
                    SET created_time = '{current_time}', updated_time = '{current_time}' 
                    WHERE created_time IS NULL OR updated_time IS NULL
                """))
                
                # 6. 更新现有submission_details的is_correct字段
                logger.info("更新现有提交详情的正确性标记...")
                
                conn.execute(text("""
                    UPDATE submission_details 
                    SET is_correct = CASE 
                        WHEN sd.selected_answer = q.answer THEN 1 
                        ELSE 0 
                    END
                    FROM submission_details sd
                    JOIN questions q ON sd.question_id = q.id
                    WHERE submission_details.id = sd.id AND submission_details.is_correct = 0
                """))
                
                # 提交事务
                trans.commit()
                logger.info("数据库优化迁移完成！")
                
            except Exception as e:
                # 回滚事务
                trans.rollback()
                logger.error(f"迁移过程中发生错误，已回滚: {e}")
                raise
                
    except Exception as e:
        logger.error(f"数据库迁移失败: {e}")
        raise

def verify_migration():
    """验证迁移结果"""
    try:
        engine = create_engine(Config.DATABASE_URL)
        
        with engine.connect() as conn:
            # 检查新字段是否存在
            logger.info("验证迁移结果...")
            
            # 检查Questions表
            result = conn.execute(text("SELECT COUNT(*) FROM pragma_table_info('questions') WHERE name IN ('difficulty_level', 'topic_category', 'created_time', 'updated_time')"))
            questions_fields = result.fetchone()[0]
            
            # 检查Submissions表
            result = conn.execute(text("SELECT COUNT(*) FROM pragma_table_info('submissions') WHERE name IN ('total_questions', 'correct_answers', 'completion_time')"))
            submissions_fields = result.fetchone()[0]
            
            # 检查SubmissionDetails表
            result = conn.execute(text("SELECT COUNT(*) FROM pragma_table_info('submission_details') WHERE name IN ('is_correct', 'answer_time')"))
            details_fields = result.fetchone()[0]
            
            # 检查索引
            result = conn.execute(text("SELECT COUNT(*) FROM sqlite_master WHERE type='index' AND name LIKE 'idx_%'"))
            index_count = result.fetchone()[0]
            
            logger.info(f"验证结果:")
            logger.info(f"  Questions表新字段: {questions_fields}/4")
            logger.info(f"  Submissions表新字段: {submissions_fields}/3")
            logger.info(f"  SubmissionDetails表新字段: {details_fields}/2")
            logger.info(f"  创建的索引数量: {index_count}")
            
            if questions_fields == 4 and submissions_fields == 3 and details_fields == 2:
                logger.info("✅ 数据库迁移验证成功！")
                return True
            else:
                logger.error("❌ 数据库迁移验证失败！")
                return False
                
    except Exception as e:
        logger.error(f"验证迁移结果失败: {e}")
        return False

if __name__ == "__main__":
    try:
        print("开始数据库优化迁移...")
        migrate_database()
        
        print("\n验证迁移结果...")
        if verify_migration():
            print("\n🎉 数据库优化迁移成功完成！")
        else:
            print("\n⚠️ 迁移可能存在问题，请检查日志")
            
    except Exception as e:
        print(f"\n❌ 迁移失败: {e}")
        sys.exit(1)